/*
 * Copyright 2024-2024 the original author or authors.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      https://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package com.mdframe.ai.agent.client;

import com.alibaba.cloud.ai.dashscope.chat.DashScopeChatOptions;
import com.mdframe.ai.agent.tools.GetTableSchemaFunction;
import com.mdframe.ai.agent.tools.LoggingAdvisor;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.PromptChatMemoryAdvisor;
import org.springframework.ai.chat.client.advisor.QuestionAnswerAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import reactor.core.publisher.Flux;

import java.time.LocalDate;

import static org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor.CHAT_MEMORY_CONVERSATION_ID_KEY;
import static org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor.CHAT_MEMORY_RETRIEVE_SIZE_KEY;

/**
 * * @author Christian Tzolov
 */
@Service
public class CustomerSupportAssistant {
    
    private final ChatClient chatClient;
    
    private final ChatClient codeGenChatClient;
    
    public CustomerSupportAssistant(ChatClient.Builder modelBuilder, VectorStore vectorStore, ChatMemory chatMemory) {
        this.chatClient = modelBuilder.defaultSystem("""
                        	你是一个代码生成助手，你可以根据数据库表并参考提供的代码示例生成代码。按照提供的Mysql表结构生成entity,mapper,service,serviceImpl,Controller的代码。
                        	不需要输出表结构，只需要生成代码。生成之后的代码调用 ToFile 保存到本地文件，如果用户不需要生成代码，你可以正常回答用户的问题
                        """).defaultAdvisors(new PromptChatMemoryAdvisor(chatMemory), // Chat Memory
                        // new VectorStoreChatMemoryAdvisor(vectorStore)),
                        new LoggingAdvisor(), // Logging
                        new QuestionAnswerAdvisor(vectorStore, SearchRequest.defaults()))
                .defaultFunctions("getTableSchema", "ToFile")
                //.defaultFunction("getTableSchema","根据数据库表生成代码",new GetTableSchemaFunction())
                .build();
        this.codeGenChatClient = modelBuilder.defaultSystem("""
                        	你是一个代码生成助手 根据用户提出的需求文档为用户提供完整的代码实现， 不需要参考本地的向量存储，生成之后的代码调用 ToFile 保存到本地文件
                        """).defaultAdvisors(new PromptChatMemoryAdvisor(chatMemory), // Chat Memory
                        new LoggingAdvisor(), // Logging
                        new QuestionAnswerAdvisor(vectorStore, SearchRequest.defaults()))
                .defaultFunctions("getTableSchema", "ToFile")
                .build();
    }
    
    public Flux<String> chat(String chatId, String userMessageContent) {
        return this.chatClient.prompt().system(s -> s.param("current_date", LocalDate.now().toString()))
                .system(s -> s.param("file_path", "/var/tmp/")).user(userMessageContent).advisors(
                        a -> a.param(CHAT_MEMORY_CONVERSATION_ID_KEY, chatId).param(CHAT_MEMORY_RETRIEVE_SIZE_KEY, 100))
                .stream().content();
    }
    
    public Flux<String> codeGenChat(String chatId, String userMessageContent) {
        return this.codeGenChatClient.prompt().system(s -> s.param("current_date", LocalDate.now().toString()))
                .system(s -> s.param("file_path", "/var/tmp/")).user(userMessageContent).advisors(
                        a -> a.param(CHAT_MEMORY_CONVERSATION_ID_KEY, chatId).param(CHAT_MEMORY_RETRIEVE_SIZE_KEY, 100))
                .options(
                        DashScopeChatOptions.builder().withEnableSearch(true)
                        .build())
                .stream().content();
    }
    
}
